from PyQt6.QtWidgets import (QWidget, QLabel, QPushButton, QApplication)
from PyQt6.QtCore import Qt, QPoint, QRect, pyqtSignal
from PyQt6.QtGui import QPainter, QColor, QPen, QPixmap, QCursor, QFont
import time

class ScreenSelector(QWidget):
    finished = pyqtSignal()  # 添加完成信号
    cancel_signal = pyqtSignal()  # 新增：取消信号
    
    def __init__(self):
        super().__init__()
        self.setWindowFlags(Qt.WindowType.FramelessWindowHint | Qt.WindowType.WindowStaysOnTopHint)
        self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground)
        self.setAttribute(Qt.WidgetAttribute.WA_ShowWithoutActivating)
        self.setCursor(Qt.CursorShape.CrossCursor)
        
        # 选择区域相关
        self.begin = QPoint()
        self.end = QPoint()
        self.is_drawing = False
        self.is_moving = False
        self.is_resizing = False
        self.selection_completed = False
        self.is_confirmed = False
        
        # 调整大小相关
        self.resize_handle = None
        self.handle_size = 8
        self.resize_cursors = {
            'top_left': Qt.CursorShape.SizeFDiagCursor,
            'top': Qt.CursorShape.SizeVerCursor,
            'top_right': Qt.CursorShape.SizeBDiagCursor,
            'right': Qt.CursorShape.SizeHorCursor,
            'bottom_right': Qt.CursorShape.SizeFDiagCursor,
            'bottom': Qt.CursorShape.SizeVerCursor,
            'bottom_left': Qt.CursorShape.SizeBDiagCursor,
            'left': Qt.CursorShape.SizeHorCursor
        }
        
        # 移动相关
        self.move_start = None
        self.last_update_time = 0
        self.update_interval = 16  # 约60fps

        # 多屏幕支持：geometry为所有屏幕联合区域
        screens = QApplication.screens()
        virtual_geometry = screens[0].geometry()
        for screen in screens[1:]:
            virtual_geometry = virtual_geometry.united(screen.geometry())
        self.setGeometry(virtual_geometry)
        self.virtual_geometry = virtual_geometry  # 记录虚拟桌面区域
        self.screens = screens
        self.screen_geometries = [screen.geometry() for screen in screens]
        self.screen_ratios = [screen.devicePixelRatio() for screen in screens]

        # 遮罩：全虚拟桌面
        self.overlay = QPixmap(self.virtual_geometry.size())
        self.overlay.fill(QColor(0, 0, 0, 100))  # 降低透明度

        # 提示信息
        self.tip_label = QLabel(self)
        self.tip_label.setStyleSheet("""
            QLabel {
                color: white;
                background-color: rgba(32, 32, 32, 180);
                padding: 8px 12px;
                border-radius: 6px;
                font-family: "Microsoft YaHei UI";
                font-size: 12px;
            }
        """)
        self.tip_label.hide()
        
        # 确认按钮
        self.confirm_btn = QPushButton("√", self)
        self.confirm_btn.setStyleSheet("""
            QPushButton {
                background-color: rgba(128, 128, 128, 0.2);
                color: #808080;
                border: none;
                border-radius: 2px;
                font-size: 14px;
                padding:5px 10px;
                font-family: "Microsoft YaHei UI";
            }
            QPushButton:hover {
                background-color: rgba(128, 128, 128, 0.6);
                border-radius: 2px;
            }
            QPushButton:pressed {
                background-color: rgba(128, 128, 128, 0.4);
            }
        """)
        self.confirm_btn.clicked.connect(self.confirm_selection)
        self.confirm_btn.hide()

    def should_update(self):
        current_time = time.time() * 1000  # 转换为毫秒
        if current_time - self.last_update_time >= self.update_interval:
            self.last_update_time = current_time
            return True
        return False

    def confirm_selection(self):
        if self.selection_completed:
            self.is_confirmed = True
            self.hide()  # 先隐藏窗口，确保遮罩消失
            self.close()

    def get_selection_rect(self):
        """获取选择区域，输出为全局坐标（虚拟桌面坐标系，未做DPI换算）"""
        rect = QRect(self.begin, self.end).normalized()
        return rect

    def get_physical_rect(self):
        """获取选择区域的物理像素坐标（考虑DPI缩放）"""
        selection_rect = QRect(self.begin, self.end).normalized()
        # 只考虑选区中心点在哪个屏幕
        center = selection_rect.center()
        for i, geo in enumerate(self.screen_geometries):
            if geo.contains(center):
                ratio = self.screen_ratios[i]
                break
        else:
            ratio = 1.0
        # 换算物理像素
        physical_rect = QRect(
            int((selection_rect.x() - self.virtual_geometry.x()) * ratio),
            int((selection_rect.y() - self.virtual_geometry.y()) * ratio),
            int(selection_rect.width() * ratio),
            int(selection_rect.height() * ratio)
        )
        return physical_rect

    def get_qt_handle_rects(self, rect):
        handles = {}
        handles['top_left'] = QRect(rect.left() - self.handle_size//2, rect.top() - self.handle_size//2, self.handle_size, self.handle_size)
        handles['top'] = QRect(rect.center().x() - self.handle_size//2, rect.top() - self.handle_size//2, self.handle_size, self.handle_size)
        handles['top_right'] = QRect(rect.right() - self.handle_size//2, rect.top() - self.handle_size//2, self.handle_size, self.handle_size)
        handles['right'] = QRect(rect.right() - self.handle_size//2, rect.center().y() - self.handle_size//2, self.handle_size, self.handle_size)
        handles['bottom_right'] = QRect(rect.right() - self.handle_size//2, rect.bottom() - self.handle_size//2, self.handle_size, self.handle_size)
        handles['bottom'] = QRect(rect.center().x() - self.handle_size//2, rect.bottom() - self.handle_size//2, self.handle_size, self.handle_size)
        handles['bottom_left'] = QRect(rect.left() - self.handle_size//2, rect.bottom() - self.handle_size//2, self.handle_size, self.handle_size)
        handles['left'] = QRect(rect.left() - self.handle_size//2, rect.center().y() - self.handle_size//2, self.handle_size, self.handle_size)
        return handles

    def update_tip_label(self, pos=None):
        if not self.selection_completed:
            self.tip_label.setText("拖动鼠标选择区域")
            self.tip_label.adjustSize()
            if pos:
                self.tip_label.move(pos.x() + 10, pos.y() + 10)
            self.tip_label.show()
        else:
            rect = QRect(self.begin, self.end).normalized()
            self.tip_label.setText(f"位置: ({rect.x()}, {rect.y()}) 大小: {rect.width()}x{rect.height()} esc取消选择区域")
            self.tip_label.adjustSize()
            self.tip_label.move(self.mapFromGlobal(rect.topLeft()).x(), self.mapFromGlobal(rect.topLeft()).y() - self.tip_label.height() - 5)
            self.tip_label.show()

    def update_confirm_button(self):
        rect = QRect(self.begin, self.end).normalized()
        self.confirm_btn.adjustSize()
        button_x = self.mapFromGlobal(rect.bottomRight()).x() - self.confirm_btn.width()
        button_y = self.mapFromGlobal(rect.bottomRight()).y() - self.confirm_btn.height()
        self.confirm_btn.move(button_x, button_y)
        self.confirm_btn.show()

    def paintEvent(self, event):
        painter = QPainter(self)
        # 绘制全虚拟桌面遮罩
        painter.drawPixmap(0, 0, self.overlay)
        # 新增：绘制快捷键提示
        painter.setPen(QColor(255, 255, 255))
        painter.setFont(QFont("Microsoft YaHei UI", 16, QFont.Weight.Bold))
        text = "Esc 取消 | Ctrl+Q 快捷选区"
        margin = 40
        painter.drawText(margin, margin, text)
        if self.is_drawing or self.is_moving or self.is_resizing or self.selection_completed:
            rect = QRect(self.mapFromGlobal(self.begin), self.mapFromGlobal(self.end)).normalized()
            painter.setPen(QPen(QColor(198, 191, 188, 180), 2, Qt.PenStyle.DashLine))
            painter.drawRect(rect)
            if self.is_drawing:
                # 初次拖动时，整个选区内透明
                painter.setCompositionMode(QPainter.CompositionMode.CompositionMode_Clear)
                painter.fillRect(rect, Qt.GlobalColor.transparent)
            else:
                # 其它情况，选区内边缘5px透明，其它有遮罩
                painter.save()
                painter.setCompositionMode(QPainter.CompositionMode.CompositionMode_Clear)
                # 先清除整个选区
                # painter.fillRect(rect, Qt.GlobalColor.transparent)  # 不再全部清除
                # 清除内层区域，留下5px边缘透明
                inner_rect = rect.adjusted(5, 5, -5, -5)
                if inner_rect.width() > 0 and inner_rect.height() > 0:
                    painter.fillRect(inner_rect, Qt.GlobalColor.transparent)
                painter.restore()
            # 绘制调整手柄（仅在已完成选择时）
            if self.selection_completed:
                handles = self.get_qt_handle_rects(rect)
                for handle_rect in handles.values():
                    painter.setPen(QPen(QColor(198, 191, 188, 180), 1))
                    painter.setBrush(QColor(198, 191, 188, 180))
                    painter.drawRect(handle_rect)

    def mousePressEvent(self, event):
        if event.button() == Qt.MouseButton.LeftButton:
            global_pos = self.mapToGlobal(event.pos())
            if self.selection_completed:
                rect = QRect(self.begin, self.end).normalized()
                handles = self.get_qt_handle_rects(QRect(self.mapFromGlobal(self.begin), self.mapFromGlobal(self.end)).normalized())
                clicked_handle = None
                for handle_name, handle_rect in handles.items():
                    if handle_rect.contains(event.pos()):
                        clicked_handle = handle_name
                        break
                if clicked_handle:
                    self.is_resizing = True
                    self.resize_handle = clicked_handle
                    self.resize_start = self.mapToGlobal(event.pos())
                    self.original_rect = QRect(self.begin, self.end)
                else:
                    expanded_rect = QRect(
                        self.mapFromGlobal(QPoint(rect.x() - 10, rect.y() - 10)),
                        self.mapFromGlobal(QPoint(rect.bottomRight().x() + 10, rect.bottomRight().y() + 10))
                    )
                    if expanded_rect.contains(event.pos()):
                        self.is_moving = True
                        self.move_start = self.mapToGlobal(event.pos())
            else:
                self.begin = global_pos
                self.end = global_pos
                self.is_drawing = True

    def mouseMoveEvent(self, event):
        global_pos = self.mapToGlobal(event.pos())
        if self.is_drawing:
            self.end = global_pos
            if self.should_update():
                self.update_tip_label(event.pos())
                self.update()
        elif self.is_moving:
            diff = global_pos - self.move_start
            self.begin += diff
            self.end += diff
            self.move_start = global_pos
            if self.should_update():
                self.update_tip_label()
                self.update_confirm_button()
                self.update()
        elif self.is_resizing:
            rect = self.original_rect
            new_rect = QRect(rect)
            offset = global_pos - self.resize_start
            if 'right' in self.resize_handle:
                new_rect.setRight(rect.right() + offset.x())
            if 'left' in self.resize_handle:
                new_rect.setLeft(rect.left() + offset.x())
            if 'bottom' in self.resize_handle:
                new_rect.setBottom(rect.bottom() + offset.y())
            if 'top' in self.resize_handle:
                new_rect.setTop(rect.top() + offset.y())
            self.begin = new_rect.topLeft()
            self.end = new_rect.bottomRight()
            if self.should_update():
                self.update_tip_label()
                self.update_confirm_button()
                self.update()
        elif self.selection_completed:
            rect = QRect(self.mapFromGlobal(self.begin), self.mapFromGlobal(self.end)).normalized()
            handles = self.get_qt_handle_rects(rect)
            cursor_handle = None
            for handle_name, handle_rect in handles.items():
                if handle_rect.contains(event.pos()):
                    cursor_handle = handle_name
                    break
            if cursor_handle:
                self.setCursor(self.resize_cursors[cursor_handle])
            else:
                expanded_rect = QRect(
                    rect.x() - 10,
                    rect.y() - 10,
                    rect.width() + 20,
                    rect.height() + 20
                )
                if expanded_rect.contains(event.pos()):
                    self.setCursor(Qt.CursorShape.SizeAllCursor)
                else:
                    self.setCursor(Qt.CursorShape.ArrowCursor)

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.MouseButton.LeftButton:
            if self.is_drawing:
                self.is_drawing = False
                self.selection_completed = True
                self.update_tip_label()
                QApplication.processEvents()
                self.update_confirm_button()
                self.update()
            elif self.is_moving:
                self.is_moving = False
            elif self.is_resizing:
                self.is_resizing = False
                self.resize_handle = None

    def keyPressEvent(self, event):
        if event.key() == Qt.Key.Key_Escape:
            self.cancel_signal.emit()  # 新增：发射取消信号
            self.close()
        elif event.key() == Qt.Key.Key_Return or event.key() == Qt.Key.Key_Enter:
            if self.selection_completed:
                self.confirm_selection()
        elif self.selection_completed:
            step = 1 if not event.modifiers() & Qt.KeyboardModifier.ShiftModifier else 10
            if event.key() == Qt.Key.Key_Left:
                self.begin.setX(self.begin.x() - step)
                self.end.setX(self.end.x() - step)
            elif event.key() == Qt.Key.Key_Right:
                self.begin.setX(self.begin.x() + step)
                self.end.setX(self.end.x() + step)
            elif event.key() == Qt.Key.Key_Up:
                self.begin.setY(self.begin.y() - step)
                self.end.setY(self.end.y() - step)
            elif event.key() == Qt.Key.Key_Down:
                self.begin.setY(self.begin.y() + step)
                self.end.setY(self.end.y() + step)

    def closeEvent(self, event):
        self.finished.emit()
        event.accept()

    def showEvent(self, event):
        super().showEvent(event)
        self.raise_()
        self.activateWindow() 